package workflow

import (
	"context"
	"fmt"
	"sync"
	"time"

	"github.com/sirupsen/logrus"

	"xagent/internal/actors"
	proto "xagent/proto/v1"
)

// ExecutionResult represents the result of a task execution
type ExecutionResult struct {
	TaskID    string
	Success   bool
	Output    string
	Error     error
	StartTime time.Time
	EndTime   time.Time
	Metrics   map[string]interface{}
}

// TaskExecutor handles task execution in workflows
type TaskExecutor struct {
	workflow   *SmartWorkflow
	results    map[string]*ExecutionResult
	agents     map[string]*actors.AIAgent
	maxRetries int
	timeout    time.Duration
	mutex      sync.RWMutex
	logger     *logrus.Logger
}

// NewTaskExecutor creates a new task executor
func NewTaskExecutor(workflow *SmartWorkflow) *TaskExecutor {
	return &TaskExecutor{
		workflow:   workflow,
		results:    make(map[string]*ExecutionResult),
		agents:     make(map[string]*actors.AIAgent),
		maxRetries: 3,
		timeout:    30 * time.Minute,
		logger:     logrus.New(),
	}
}

// RegisterAgent registers an AI agent for a specific role
func (e *TaskExecutor) RegisterAgent(role string, agent *actors.AIAgent) {
	e.mutex.Lock()
	defer e.mutex.Unlock()
	e.agents[role] = agent
}

// ExecuteStage executes all tasks in a workflow stage
func (e *TaskExecutor) ExecuteStage(ctx context.Context, stage *SmartWorkflowStage) error {
	e.logger.WithFields(logrus.Fields{
		"stage": stage.Name,
		"role":  stage.Role,
	}).Info("Executing stage")

	// Get agent for the stage
	agent, ok := e.agents[stage.Role]
	if !ok {
		return fmt.Errorf("no agent found for role: %s", stage.Role)
	}

	// Execute each action in the stage
	for _, action := range stage.Actions {
		task := &proto.Task{
			Id:          fmt.Sprintf("%s-%s", stage.Name, action),
			Name:        action,
			Description: fmt.Sprintf("Execute %s for stage %s", action, stage.Name),
			Status:      proto.TaskStatus_TASK_STATUS_PENDING,
			Metadata: map[string]string{
				"stage": stage.Name,
				"role":  stage.Role,
			},
		}

		result, err := e.executeTask(ctx, agent, task)
		if err != nil {
			e.logger.WithError(err).Error("Task execution failed")
			return err
		}

		e.mutex.Lock()
		e.results[task.Id] = result
		e.mutex.Unlock()

		if !result.Success {
			return fmt.Errorf("task execution failed: %v", result.Error)
		}
	}

	return nil
}

// executeTask executes a single task with retries
func (e *TaskExecutor) executeTask(ctx context.Context, agent *actors.AIAgent, task *proto.Task) (*ExecutionResult, error) {
	var lastErr error
	startTime := time.Now()

	for attempt := 0; attempt < e.maxRetries; attempt++ {
		// Create task adapter
		taskMsg := NewTaskAdapter(task)

		// Execute task
		err := agent.ProcessTask(taskMsg.task)
		if err != nil {
			lastErr = err
			continue
		}

		// Record successful execution
		result := &ExecutionResult{
			TaskID:    task.Id,
			Success:   true,
			StartTime: startTime,
			EndTime:   time.Now(),
			Metrics: map[string]interface{}{
				"duration": time.Now().Sub(startTime).String(),
				"attempts": attempt + 1,
			},
		}

		e.mutex.Lock()
		e.results[task.Id] = result
		e.mutex.Unlock()

		// Task succeeded
		return result, nil
	}

	return nil, fmt.Errorf("task execution failed after %d attempts: %v", e.maxRetries, lastErr)
}

// GetTaskResult returns the result of a specific task
func (e *TaskExecutor) GetTaskResult(taskID string) (*ExecutionResult, bool) {
	e.mutex.RLock()
	defer e.mutex.RUnlock()
	result, exists := e.results[taskID]
	return result, exists
}

// GetStageResults returns all task results for a stage
func (e *TaskExecutor) GetStageResults(stageName string) []*ExecutionResult {
	e.mutex.RLock()
	defer e.mutex.RUnlock()

	var results []*ExecutionResult
	for _, result := range e.results {
		if result.TaskID[:len(stageName)] == stageName {
			results = append(results, result)
		}
	}
	return results
}

// GetExecutionMetrics returns execution metrics for all tasks
func (e *TaskExecutor) GetExecutionMetrics() map[string]interface{} {
	e.mutex.RLock()
	defer e.mutex.RUnlock()

	metrics := make(map[string]interface{})
	var totalDuration time.Duration
	successCount := 0
	failureCount := 0

	for _, result := range e.results {
		if result.Success {
			successCount++
		} else {
			failureCount++
		}
		totalDuration += result.EndTime.Sub(result.StartTime)
	}

	metrics["total_tasks"] = len(e.results)
	metrics["successful_tasks"] = successCount
	metrics["failed_tasks"] = failureCount
	metrics["success_rate"] = float64(successCount) / float64(len(e.results))
	metrics["average_duration"] = totalDuration / time.Duration(len(e.results))

	return metrics
}

// SetTimeout sets the execution timeout for tasks
func (e *TaskExecutor) SetTimeout(timeout time.Duration) {
	e.timeout = timeout
}

// SetMaxRetries sets the maximum number of retries for failed tasks
func (e *TaskExecutor) SetMaxRetries(retries int) {
	e.maxRetries = retries
}
